// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/functions/prestosql/aggregates/count_if_aggregate.h>
#include <pollux/exec/aggregate.h>
#include <pollux/expression/function_signature.h>
#include <pollux/functions/prestosql/aggregates/aggregate_names.h>
#include <pollux/vector/decoded_vector.h>
#include <pollux/vector/flat_vector.h>
#include <pollux/vector/simple_vector.h>

namespace kumo::pollux::aggregate::prestosql {
    namespace {
        class CountIfAggregate : public exec::Aggregate {
        public:
            explicit CountIfAggregate() : exec::Aggregate(BIGINT()) {
            }

            int32_t accumulatorFixedWidthSize() const override {
                return sizeof(int64_t);
            }

            void extractAccumulators(char **groups, int32_t numGroups, VectorPtr *result)
            override {
                extractValues(groups, numGroups, result);
            }

            void extractValues(char **groups, int32_t numGroups, VectorPtr *result)
            override {
                auto *vector = (*result)->as<FlatVector<int64_t> >();
                POLLUX_CHECK(vector);
                vector->resize(numGroups);

                auto *rawValues = vector->mutableRawValues();
                for (vector_size_t i = 0; i < numGroups; ++i) {
                    rawValues[i] = *value<int64_t>(groups[i]);
                }
            }

            void addRawInput(
                char **groups,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*mayPushdown*/) override {
                DecodedVector decoded(*args[0], rows);

                if (decoded.isConstantMapping()) {
                    if (decoded.is_null_at(0)) {
                        return;
                    }
                    if (decoded.value_at<bool>(0)) {
                        rows.applyToSelected(
                            [&](vector_size_t i) { addToGroup(groups[i], 1); });
                    }
                } else if (decoded.may_have_nulls()) {
                    rows.applyToSelected([&](vector_size_t i) {
                        if (decoded.is_null_at(i)) {
                            return;
                        }
                        if (decoded.value_at<bool>(i)) {
                            addToGroup(groups[i], 1);
                        }
                    });
                } else {
                    rows.applyToSelected([&](vector_size_t i) {
                        if (decoded.value_at<bool>(i)) {
                            addToGroup(groups[i], 1);
                        }
                    });
                }
            }

            void addIntermediateResults(
                char **groups,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*mayPushdown*/) override {
                DecodedVector decoded(*args[0], rows);

                if (decoded.isConstantMapping()) {
                    auto numTrue = decoded.value_at<int64_t>(0);
                    rows.applyToSelected(
                        [&](vector_size_t i) { addToGroup(groups[i], numTrue); });
                    return;
                }

                rows.applyToSelected([&](vector_size_t i) {
                    auto numTrue = decoded.value_at<int64_t>(i);
                    addToGroup(groups[i], numTrue);
                });
            }

            void addSingleGroupRawInput(
                char *group,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*mayPushdown*/) override {
                DecodedVector decoded(*args[0], rows);

                // Constant mapping - check once and add number of selected rows if true.
                if (decoded.isConstantMapping()) {
                    if (!decoded.is_null_at(0)) {
                        auto isTrue = decoded.value_at<bool>(0);
                        if (isTrue) {
                            addToGroup(group, rows.countSelected());
                        }
                    }
                    return;
                }

                int64_t numTrue = 0;
                if (decoded.may_have_nulls()) {
                    rows.applyToSelected([&](vector_size_t i) {
                        if (decoded.is_null_at(i)) {
                            return;
                        }
                        if (decoded.value_at<bool>(i)) {
                            ++numTrue;
                        }
                    });
                } else {
                    rows.applyToSelected([&](vector_size_t i) {
                        if (decoded.value_at<bool>(i)) {
                            ++numTrue;
                        }
                    });
                }
                addToGroup(group, numTrue);
            }

            void addSingleGroupIntermediateResults(
                char *group,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*mayPushdown*/) override {
                auto arg = args[0]->as<SimpleVector<int64_t> >();

                int64_t numTrue = 0;
                rows.applyToSelected([&](auto row) { numTrue += arg->value_at(row); });

                addToGroup(group, numTrue);
            }

        protected:
            void initializeNewGroupsInternal(
                char **groups,
                melon::Range<const vector_size_t *> indices) override {
                for (auto i: indices) {
                    *value<int64_t>(groups[i]) = 0;
                }
            }

        private:
            inline void addToGroup(char *group, int64_t numTrue) {
                *value<int64_t>(group) += numTrue;
            }
        };
    } // namespace

    void registerCountIfAggregate(
        const std::string &prefix,
        bool withCompanionFunctions,
        bool overwrite) {
        std::vector<std::shared_ptr<exec::AggregateFunctionSignature> > signatures{
            exec::AggregateFunctionSignatureBuilder()
            .returnType("bigint")
            .intermediateType("bigint")
            .argumentType("boolean")
            .build(),
        };

        auto name = prefix + kCountIf;
        exec::registerAggregateFunction(
            name,
            std::move(signatures),
            [name](
        core::AggregationNode::Step step,
        std::vector<TypePtr> argTypes,
        const TypePtr & /*resultType*/,
        const core::QueryConfig & /*config*/)
        -> std::unique_ptr<exec::Aggregate> {
                POLLUX_CHECK_EQ(argTypes.size(), 1, "{} takes one argument", name);

                auto isPartial = exec::isRawInput(step);
                if (isPartial) {
                    POLLUX_CHECK_EQ(
                        argTypes[0]->kind(),
                        TypeKind::BOOLEAN,
                        "{} function only accepts boolean parameter",
                        name);
                }

                return std::make_unique<CountIfAggregate>();
            },
            {false /*orderSensitive*/, false /*companionFunction*/},
            withCompanionFunctions,
            overwrite);
    }
} // namespace kumo::pollux::aggregate::prestosql
